import cv2
import numpy as np
import os

"缩放"
# 缩放
def Scale(image, scale):
    return cv2.resize(image, None, fx=scale, fy=scale)


"反转"
# 水平翻转
def Horizontal(image):
    return cv2.flip(image,1,dst=None)

# 垂直反转
def Vertical(image):
    return cv2.flip(image,0,dst=None)

# 旋转，R可控制图片放大缩小
def Rotate(image, angle=15, scale=0.9):
    w = image.shape[1]
    h = image.shape[0]
    
    # 旋转矩阵
    M = cv2.getRotationMatrix2D((w/2, h/2), angle, scale)
    # 旋转
    image = cv2.warpAffine(image,M,(w,h))
    return image

"明亮度"
# 变暗
def Darker(image, percentage=0.5):
    image_copy = image.copy()
    w = image.shape[1]
    h = image.shape[0]
    
    # get darker
    for xi in range(0,w):
        for xj in range(0,h):
            image_copy[xj,xi,0] = int(image[xj,xi,0]*percentage)
            image_copy[xj,xi,1] = int(image[xj,xi,1]*percentage)
            image_copy[xj,xi,2] = int(image[xj,xi,2]*percentage)
    
    return image_copy
            
# 变亮
def Brighter(image, percentage=2):
    image_copy = image.copy()
    w = image.shape[1]
    h = image.shape[0]
    
    # get brighter
    for xi in range(0,w):
        for xj in range(0,h):
            image_copy[xj,xi,0] = np.clip(int(image[xj,xi,0]*percentage), a_max=255, a_min=0)
            image_copy[xj,xi,1] = np.clip(int(image[xj,xi,1]*percentage), a_max=255, a_min=0)
            image_copy[xj,xi,2] = np.clip(int(image[xj,xi,2]*percentage), a_max=255, a_min=0)
    return image_copy
        
# 平移
def Move(img, x, y):
    img_info = img.shape
    height = img_info[0]
    width = img_info[1]
    
    # x表示向上平移距离，y表示向右平移距离，其他部分表示为单位矩阵，所以变换前后不变
    mat_translation = np.float32([[1,0,x],[0,1,y]])
    dst = cv2.warpAffine(img, mat_translation, (width, height)) 
    return dst   

"增加噪声"
# 椒盐噪声
def SaltAndPepper(src,percetage=0.05):
    SP_NoiseImg = src.copy()
    SP_NoiseNum = int(percetage*src.shape[0]-1)
    for i in range(SP_NoiseNum):
        randR = np.random.randint(0,src.shape[0]-1)
        randG = np.random.randint(0,src.shape[1]-1)
        randB = np.random.randint(0,3)
        if np.random.randint(0,1)==0:
            SP_NoiseImg[randR, randG, randB]=0
        else:
            SP_NoiseImg[randR, randG, randB]=255
    return SP_NoiseImg


# 高斯噪声
def GaussianNoise(image, percetage=0.05):
    G_Noiseimg = image.copy()
    w = image.shape[1]
    h = image.shape[0]
    G_NoiseeNum = int(percetage*image.shape[0]*image.shape[1])
    for i in range(G_NoiseeNum):
        temp_x = np.random.randint(0,h)
        temp_y = np.random.randint(0,w)
        G_NoiseeNum[temp_x][temp_y][np.random.randint(3)] = np.random.randint(1)[0]
    return G_Noiseimg

# 模糊
def Blur(img):
    blur = cv2.GaussianBlur(img, (7,7), 1.5)
    return blur
        

"对所有样本数据增强"
def TestOneDir():
    root_path = r"F:\SCIENCE_AND_MATH\Machine Learning\MCM\2021C\dataset\data_classified\Negative"
    save_path = r"F:\SCIENCE_AND_MATH\Machine Learning\MCM\2021C\dataset\data_classified\Negative"
    
    for root, dirs, files in os.walk(root_path):
        for file_i in files:
            flie_i_path = os.path.join(root, file_i)
            print(flie_i_path)
            # print(file_i)
            img_i = cv2.imread(flie_i_path)
            
            img_scale = Scale(img_i,1.5)
            cv2.imwrite(os.path.join(save_path, file_i + "_scale.jpg"), img_scale)
    
            img_horizontal = Horizontal(img_i)
            cv2.imwrite(os.path.join(save_path, file_i + "_horizontal.jpg"), img_horizontal)
            
            img_vertical = Vertical(img_i)
            cv2.imwrite(os.path.join(save_path, file_i + "_vertical.jpg"), img_vertical)
    
            img_rotate = Rotate(img_i,90)
            cv2.imwrite(os.path.join(save_path, file_i + "_rotate90.jpg"), img_rotate)
            
            img_rotate = Rotate(img_i,180)
            cv2.imwrite(os.path.join(save_path, file_i + "_rotate180.jpg"), img_rotate)
            
            img_rotate = Rotate(img_i,270)
            cv2.imwrite(os.path.join(save_path, file_i + "_rotate270.jpg"), img_rotate)
            
            img_move = Move(img_i,15,15)
            cv2.imwrite(os.path.join(save_path, file_i + "_move.jpg"), img_move)
            
            img_darker = Darker(img_i)
            cv2.imwrite(os.path.join(save_path, file_i + "_darker.jpg"), img_darker)
            
            img_brighter = Brighter(img_i)
            cv2.imwrite(os.path.join(save_path, file_i + "_brighter.jpg"), img_brighter)
            
            img_blur = Blur(img_i)
            cv2.imwrite(os.path.join(save_path, file_i + "_blur.jpg"), img_blur)
            
            img_salt = SaltAndPepper(img_i, 0.05)
            cv2.imwrite(os.path.join(save_path, file_i + "_salt.jpg"), img_salt)
            
    
if __name__=="__main__":
    TestOneDir()